import os
import pickle

import numpy as np
import matplotlib.colors as mcolors
import plotly.graph_objects as go

from matplotlib import pyplot as plt

from utils import read_mat_file, local_to_global_coords

def compute_rgb(Z, s, DN):
    hue = 5/6 * (1 - (Z - Z.min()) / (Z.max() - Z.min())) - 1/6
    Z_trimmed = Z[Z <= np.percentile(Z, 95)]  # 砍掉最大的5%
    Z_normalized = (Z - Z_trimmed.min()) / (Z_trimmed.max() - Z_trimmed.min())
    hue = 2/3 * (1 - Z_normalized)  # 1对应蓝色，0对应红色
    value =  s/s#1 - (DN - DN.min()) / (DN.max() - DN.min())
    saturation = s/s

    hue = [x if x >= 0 else 0 for x in hue]

    value =  (s - s.min()) / (s.max() - s.min())
    saturation = 1 - (DN - DN.min()) / (DN.max() - DN.min())
    
    hsv = np.stack((hue, saturation, value), axis=1)
    rgb = mcolors.hsv_to_rgb(hsv)
    return rgb

def filter_data(X, Y, Z, x_range, y_range, azi, ele):
    mask = (Z != 0) & \
        (X >= x_range[0]) & (X <= x_range[1]) & \
        (Y >= y_range[0]) & (Y <= y_range[1]) & ((azi > 0.5) & (ele > 0.85))
    return mask

lat0 = None
lon0 = None
h0 = None

min_z = 1000
max_z = -1000

def process_mat_file(mat_data, all_data, x_range, y_range):
    global lat0
    global lon0
    global h0

    start_frame_llh = mat_data.get('StartFrameLLH', None)
    start_frame_prh = mat_data.get('StartFrameRPY', None)
    azi_map = mat_data.get('aziLinearlityMap', None)
    ele_map = mat_data.get('eleLinearlityMap', None)
    target_x = mat_data.get('targetX', None)
    target_y = mat_data.get('targetY', None)
    target_z = mat_data.get('targetZ', None)
    dn = mat_data.get('DN', None)

    global max_z
    if max_z < np.max(target_z):
        max_z = np.max(target_z)

    global min_z
    if min_z > np.min(target_z):
        min_z = np.min(target_z)

    if not all([azi_map is not None, ele_map is not None, target_x is not None, 
                target_y is not None, target_z is not None, dn is not None, start_frame_llh is not None, start_frame_prh is not None]):
        raise ".mat文件中缺少数据"
    
    lat = start_frame_llh[0][0]  # 纬度
    lon = start_frame_llh[0][1]  # 精度
    h = start_frame_llh[0][2] / 1000  # 海拔高度

    pitch = start_frame_prh[0][0]  # 俯仰角（重力基准）
    roll = start_frame_prh[0][1]  # 滚转角（重力基准）
    heading = start_frame_prh[0][2]  # 与北方向夹角

    if lat0 == None:
        lat0 = lat
        lon0 = lon
        h0 = h
    X = target_x.flatten()
    Y = target_y.flatten()
    Z = target_z.flatten()
    DN = dn.flatten()
    azi = azi_map.flatten()
    ele = ele_map.flatten()

    s = azi + ele

    mask = filter_data(X, Y, Z, x_range, y_range, azi, ele)
    X = X[mask]
    Y = Y[mask]
    Z = Z[mask]
    DN = DN[mask]
    s = s[mask]

    X, Y, Z = local_to_global_coords(X, Y, Z, lat, lon, h, pitch, roll, heading, lat0, lon0, h0)

    rgb = compute_rgb(Z, s, DN)

    all_data['x'].extend(X)
    all_data['y'].extend(Y)
    all_data['z'].extend(Z)
    all_data['dn'].extend(DN)
    all_data['rgb'].extend(rgb)
    all_data['aziLinearlityMap'].extend(azi)
    all_data['eleLinearlityMap'].extend(ele)

    return all_data


def plot_combined_map(all_data, output_folder):
    X = np.array(all_data['x'])
    Y = np.array(all_data['y'])
    Z = np.array(all_data['z'])
    DN = np.array(all_data['dn'])
    rgb = np.array(all_data['rgb'])

    fig, ax = plt.subplots(figsize=(10, 10), dpi=1024//10, facecolor='black')
    ax.set_facecolor('blue')
    ax.set_axis_off()
    ax.set_aspect('equal', adjustable='datalim')
    ax.scatter(X, Y, c=rgb, s=1, marker='.')

    output_file = os.path.join(output_folder, "combined_map.png")
    fig.savefig(output_file, dpi=1024//10, facecolor=fig.get_facecolor())
    plt.close()

def plot_3d_pointcloud(all_data):
    # 从输入数据提取坐标和颜色
    X = all_data['x']
    Y = all_data['y']
    Z = all_data['z']
    rgb = all_data['rgb']

    # 创建点云
    fig = go.Figure(data=[go.Scatter3d(
        x=X, y=Y, z=Z,
        mode='markers',
        marker=dict(size=1, color=rgb, opacity=1.0)
    )])

    
    # 设置布局以保持实际尺寸比例
    fig.update_layout(scene=dict( 
        aspectmode='data',  # 使用实际数据比例
        xaxis=dict(title='X', showgrid=False),  # 关闭 X 轴网格
        yaxis=dict(title='Y', showgrid=False),  # 关闭 Y 轴网格
        zaxis=dict(title='Z', showgrid=False)   # 关闭 Z 轴网格
        ),
        template='plotly_dark',  # 使用黑色主题
        paper_bgcolor='black',   # 背景颜色设置为黑色
        plot_bgcolor='black'     # 绘图区背景颜色设置为黑色
    )

    fig.show()

def main():
    folder_path = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR"
    cache_file = os.path.join(folder_path, "all_data_cache.pkl")

    os.makedirs(output_folder, exist_ok=True)

    if os.path.exists(cache_file):
        print("加载缓存数据...")
        with open(cache_file, 'rb') as f:
            all_data = pickle.load(f)
    else:
        all_data = {'x': [], 'y': [], 'z': [], 'dn': [], 'rgb': [], 'aziLinearlityMap': [], 'eleLinearlityMap': []}

        for root, _, files in os.walk(folder_path):
            for i in range(0, len(files), 10):
                file = files[i]
                if file.endswith(".mat"):
                    file_path = os.path.join(root, file)
                    print(f"正在处理文件: {file_path}")

                    mat_data = read_mat_file(file_path)
                    if mat_data:
                        all_data = process_mat_file(
                            mat_data, 
                            all_data, 
                            x_range=(-1000, 1000),
                            y_range=(0, 90)
                        )
        
    plot_3d_pointcloud(all_data)

if __name__ == "__main__":
    main()